// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/lib/aggregates/value_list.h>
#include <pollux/exec/container_row_serde.h>

namespace kumo::pollux::aggregate {
    void ValueList::prepareAppend(HashStringAllocator *allocator) {
        if (!nullsBegin_) {
            nullsBegin_ = allocator->allocate(HashStringAllocator::kMinAlloc);
            nullsCurrent_ = {nullsBegin_, nullsBegin_->begin()};
        }

        if (!dataBegin_) {
            dataBegin_ = allocator->allocate(kInitialSize);
            dataCurrent_ = {dataBegin_, dataBegin_->begin()};
        }

        if (size_ && size_ % 64 == 0) {
            writeLastNulls(allocator);
            lastNulls_ = 0;
            // Make sure there is space for another word of null flags without
            // allocating more space. This is needed for finalize to finish in
            // constant space.
            allocator->ensureAvailable(sizeof(int64_t), nullsCurrent_);
        }
    }

    void ValueList::writeLastNulls(HashStringAllocator *allocator) {
        ByteOutputStream stream(allocator);
        if (nullsBegin_) {
            allocator->extendWrite(nullsCurrent_, stream);
        } else {
            auto position = allocator->newWrite(stream, kInitialSize);
            nullsBegin_ = position.header;
        }
        stream.appendOne(lastNulls_);
        nullsCurrent_ = allocator->finishWrite(stream, kInitialSize).second;
    }

    void ValueList::appendNull(HashStringAllocator *allocator) {
        prepareAppend(allocator);
        lastNulls_ |= 1UL << (size_ % 64);
        ++size_;
    }

    void ValueList::appendNonNull(
        const BaseVector &values,
        vector_size_t index,
        HashStringAllocator *allocator) {
        prepareAppend(allocator);
        ByteOutputStream stream(allocator);
        allocator->extendWrite(dataCurrent_, stream);
        // The stream may have a tail of a previous write.
        const auto initialSize = stream.size();
        static const exec::ContainerRowSerdeOptions options{};
        exec::ContainerRowSerde::serialize(values, index, stream, options);
        ++size_;
        bytes_ += stream.size() - initialSize;

        // Leave space up to half the size appended so far, at least 24 but no more
        // than 1024.
        dataCurrent_ =
                allocator->finishWrite(stream, std::clamp(bytes_ / 2, 24, 1024)).second;
    }

    void ValueList::appendValue(
        const DecodedVector &decoded,
        vector_size_t index,
        HashStringAllocator *allocator) {
        auto &base = *decoded.base();
        if (decoded.is_null_at(index)) {
            appendNull(allocator);
        } else {
            appendNonNull(base, decoded.index(index), allocator);
        }
    }

    void ValueList::appendRange(
        const VectorPtr &vector,
        vector_size_t offset,
        vector_size_t size,
        HashStringAllocator *allocator) {
        for (auto index = offset; index < offset + size; ++index) {
            if (vector->is_null_at(index)) {
                appendNull(allocator);
            } else {
                appendNonNull(*vector, index, allocator);
            }
        }
    }

    ValueListReader::ValueListReader(ValueList &values)
        : size_{values.size()},
          lastNullsStart_{size_ % 64 == 0 ? size_ - 64 : size_ - size_ % 64},
          lastNulls_{values.lastNulls()},
          dataStream_{values.dataBegin()},
          nullsStream_{values.nullsBegin()} {
    }

    bool ValueListReader::next(BaseVector &output, vector_size_t outputIndex) {
        if (pos_ == lastNullsStart_) {
            nulls_ = lastNulls_;
        } else if (pos_ % 64 == 0) {
            nulls_ = nullsStream_.read<uint64_t>();
        }

        if (nulls_ & (1UL << (pos_ % 64))) {
            output.set_null(outputIndex, true);
        } else {
            exec::ContainerRowSerde::deserialize(dataStream_, outputIndex, &output);
        }

        pos_++;
        return pos_ < size_;
    }
} // namespace kumo::pollux::aggregate
